package edu.qdu.recommand.offline

import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.ALS
import org.apache.spark.sql.SparkSession

/**
 * ALS的案例
 */

//封装数据
case class Ratings(userId:Int,movieId:Int,rating:Int)

object ALSTest {
  def main(args: Array[String]): Unit = {
    //获取sparkSQL的上下文
    val spark = SparkSession.builder()
      .appName("dataloader")
      .master("local[*]") //如果往服务器部署，需要修改一下即可
      .getOrCreate()

    //初始化数据
    import spark.implicits._
    val ratings = spark.sparkContext
      .textFile("D:/bigdata_projects/hn-recommender/src/main/resources/sample_movielens_ratings.txt")
      .map(line => {
        val fileds = line.split("::")

        //封装ratings
        Ratings(fileds(0).trim.toInt, fileds(1).trim.toInt, fileds(2).trim.toInt)
      })
      .toDF()
    ratings.show()

    //划分训练和测试数据集
    val Array(train,test) = ratings.randomSplit(Array(0.8, 0.2))
    //构建ALS模型
    val als = new ALS()
      .setRank(10) //隐士因子，其实就是维度数量
      .setMaxIter(10) //迭代次数
      .setRegParam(0.01) //正则系数
      .setColdStartStrategy("drop") //冷启动策略
      .setUserCol("userId")
      .setItemCol("movieId")
      .setRatingCol("rating")
    //训练数据
    val model = als.fit(train)
    //测试模型 --- 使用已经训练好的模型进行预测
    val prediction = model.transform(test)
    //构建一个计算器
    val evaluator = new RegressionEvaluator()
      .setMetricName("rmse")
      .setLabelCol("rating")
      .setPredictionCol("prediction")
    //使用计算器对预测数据集进行均方根误差计算
    val rmse = evaluator.evaluate(prediction)

    println(s"均方根误差为: $rmse")
    //数据推荐
    //为每个用户推荐10个电影
    val forUserMoivesDF = model.recommendForAllUsers(10)
    forUserMoivesDF.show()
    //为每个电影推荐10个用户
    val forMovieUsersDF = model.recommendForAllItems(10)
    forMovieUsersDF.show()
    //为3个用户推荐3个电影
    val frame = model.recommendForUserSubset(ratings.select(als.getUserCol).distinct().limit(3), 3)
    frame.show()
    //为3个电影推荐3个用户
    val frame1 = model.recommendForItemSubset(ratings.select(als.getItemCol).limit(3), 3)
    frame1.show()

    //关闭spark
    spark.stop()
  }
}